# -*- coding: utf-8 -*-
import numpy as np


def main(R, Q, C, F, maxIterCount, p1, p2):
    # 当前迭代次数
    count = 0
    # 回路个数
    N = len(C)
    # 解算成功或失败
    ret = False
    while not ret and count < maxIterCount:
        # 记录所有回路的修正风量和风压差
        DQ, DH = [], []
        for i in range(N):
            # h-回路阻力代数和, f-风机动力和
            h, f = 0, 0
            # hh-h求导, ff-f求导
            hh, ff = 0, 0
            # 计算回路C[i]的修正风量
            for x in C[i]:
                # j-分支索引, c方向系数
                j, c = x
                # 分支的风阻r、当前风量q
                r, q = R[j], Q[j]
                h += c * r * q * abs(q)
                hh += abs(2 * c * r * q)
                # 如果分支上有风机
                if F[j] != None:
                    a0, a1, a2 = F[j]
                    f += c * (a0 + a1 * q + a2 * q ** 2)
                    ff += c * (a1 + 2 * a2 * q)
            # 计算回路C[i]修正风量
            # 加1e-6的目的: 避免分母为0
            deltaH = h - f
            deltaQ = -1 * deltaH / (hh - ff + 1e-6)
            DQ.append(deltaQ)
            DH.append(deltaH)
            print("第%d次迭代, 第%d个回路" % (count + 1, i + 1), deltaQ, deltaH)
            # 修正回路中所有分支的风量
            for x in C[i]:
                # j-分支索引, c方向系数
                j, c = x
                Q[j] += c * deltaQ
        # 判断是否达到误差精度要求
        if np.max(np.abs(DQ)) < p1 and np.max(np.abs(DH)) < p2:
            print("已满足精度要求")
            ret = True
        else:
            # 迭代次数加1
            count += 1
    # 返回解算成功与否,以及实际迭代次数
    return ret, count


if __name__ == "__main__":
    # q1 ~ q12
    Q = [40, 20, 20, 10, 5, 5, 15, 15, 10, 30, 40, 40]
    # r1 ~ r12
    R = [
        0.055,
        0.075,
        0.16,
        0.25,
        0.16,
        0.375,
        0.18,
        0.135,
        0.22,
        0.17,
        0.085,
        0,
    ]
    # 回路列表
    C = [
        [(3, 1), (1, 1), (2, -1)],  # C1 = {e4, e2, e3}, 第一个就是基准分支，养成好习惯
        [(5, 1), (1, 1), (2, -1), (7, -1)],  # C2 = {e6, e2, e3, e8}
        [(6, 1), (1, -1), (2, 1), (4, -1)],  # C3 = {e7, e2, e3, e5}
        [
            (8, 1),
            (1, 1),
            (2, -1),
            (4, 1),
            (7, -1),
            (9, -1),
        ],  # C4 = {e9, e2, e3, e5, e8, e10}
        [
            (11, 1),
            (0, 1),
            (2, 1),
            (7, 1),
            (9, 1),
            (10, 1),
        ],  # C5 = {e12, e1, e3, e8, e10, e11}
    ]
    # 风机
    F = [
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        (1035.92, 51.73, -0.43),  # a0, a1, a2
        None,
    ]
    # 最大迭代次数
    maxIterCount = 100
    # 风量误差精度、风压误差精度
    p1, p2 = 0.001, 0.1
    ret, count = main(R, Q, C, F, maxIterCount, p1, p2)
    print("迭代结果:", ret, "实际迭代次数:", count)
    N = len(Q)
    str_list = ["r(%d)=%.4f\tq(%d)=%.2f" % (i + 1, R[i], i + 1, Q[i]) for i in range(N)]
    print("\n".join(str_list))
